package goroutine

import (
	"context"
	"sync"
	"time"

	"github.com/pkg/errors"
)

const (
	maxCurrentLoadTaskNum = 10               // 当前能够同时加载最大任务数量
	taskExecTimeout       = time.Second * 15 // 任务执行超时时间
)

type errGroup struct {
	// 非堵塞任务参数
	wg      sync.WaitGroup
	errOnce sync.Once

	// 堵塞任务参数
	blockTaskArgs []*taskArg

	ctx    context.Context
	cancel func()

	err error
}

type taskArg struct {
	callback func(args ...interface{}) error
	args     []interface{}
}

func NewErrGroup(ctxs ...context.Context) (*errGroup, context.Context) {
	var ctx context.Context
	if len(ctxs) > 0 {
		ctx = ctxs[0]
	} else {
		ctx = context.Background()
	}

	ctx, cancel := context.WithCancel(ctx)
	return &errGroup{cancel: cancel, ctx: ctx}, ctx
}

func (g *errGroup) Wait() error {
	if len(g.blockTaskArgs) > 0 {
		return g.doBlockTask()
	} else {
		return g.doTask()
	}
}

// 非堵塞性任务
func (g *errGroup) Go(f func(args ...interface{}) error, args ...interface{}) {
	g.wg.Add(1)
	go func() {
		defer g.wg.Done()

		if err := f(args...); err != nil {
			g.errOnce.Do(func() {
				g.err = err
				if g.cancel != nil {
					g.cancel()
				}
			})
		}
	}()
}

// 堵塞性任务:遇到error后续协程不执行
func (g *errGroup) GoWithBlock(f func(args ...interface{}) error, args ...interface{}) {
	g.blockTaskArgs = append(g.blockTaskArgs, &taskArg{
		callback: f,
		args:     args,
	})
}

// 执行堵塞任务
func (g *errGroup) doBlockTask() error {
	defer func() {
		if g.cancel != nil {
			g.cancel()
		}
	}()

	var (
		doCount   int // 执行次数
		errChan   = make(chan error, len(g.blockTaskArgs))
		limitChan = make(chan struct{}, maxCurrentLoadTaskNum) // 限流
	)

	for _, task := range g.blockTaskArgs {
		limitChan <- struct{}{}
		go func(limitChan chan struct{}, task *taskArg, ctx context.Context) {
			defer func() {
				<-limitChan
			}()

			errChan <- task.callback(task.args...)
		}(limitChan, task, g.ctx)
	}

	ticker := time.NewTicker(taskExecTimeout)
	defer ticker.Stop()

	for doCount < len(g.blockTaskArgs) {
		select {
		case <-ticker.C:
			return errors.New("协程执行超时")
		case err := <-errChan:
			doCount++
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// 执行无堵塞任务
func (g *errGroup) doTask() error {
	g.wg.Wait()
	if g.cancel != nil {
		g.cancel()
	}
	return g.err
}
